﻿/*
* Copyright (c) 2010-2012 Tesla Engine Group
* 
* Permission is hereby granted, free of charge, to any person obtaining a copy
* of this software and associated documentation files (the "Software"), to deal
* in the Software without restriction, including without limitation the rights
* to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
* copies of the Software, and to permit persons to whom the Software is
* furnished to do so, subject to the following conditions:
* 
* The above copyright notice and this permission notice shall be included in
* all copies or substantial portions of the Software.
* 
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
* THE SOFTWARE.
*/


using System;
using Tesla.Direct3D10.Graphics.Implementation;
using Tesla.Graphics;
using D3D = SlimDX.Direct3D10;
using D3DC = SlimDX.D3DCompiler;

namespace Tesla.Direct3D10.Graphics {
    /// <summary>
    /// Concrete Direct3D10 implementation of <see cref="IEffectPass"/>.
    /// </summary>
    public sealed class D3D10EffectPass : IEffectPass {
        private String _name;
        private D3D10EffectAnnotationCollection _annotations;
        private InputLayoutMap _layoutMap;
        private D3D10Renderer _renderer;
        private D3D.EffectPass _pass;

        /// <summary>
        /// Gets the name of the pass.
        /// </summary>
        public String Name {
            get {
                return _name;
            }
        }

        /// <summary>
        /// Gets the D3D10Renderer.
        /// </summary>
        internal D3D10Renderer Renderer {
            get {
                return _renderer;
            }
        }

        /// <summary>
        /// Gets the input layout associated with this pass, and other passes from the same shader source.
        /// </summary>
        internal InputLayoutMap InputLayoutMap {
            get {
                return _layoutMap;
            }
        }

        /// <summary>
        /// Gets the collection of annotations that decorates the pass.
        /// </summary>
        public IEffectAnnotationCollection Annotations {
            get {
                return _annotations;
            }
        }

        /// <summary>
        /// Creates a new instance of <see cref="D3D10EffectPass"/>.
        /// </summary>
        /// <param name="renderer">The D3D10 renderer.</param>
        /// <param name="pass">The D3D10 pass.</param>
        /// <param name="map">The input layout map to share between passes that have the same effect origin.</param>
        internal D3D10EffectPass(D3D10Renderer renderer, D3D.EffectPass pass, InputLayoutMap map) {
            _pass = pass;
            _renderer = renderer;

            D3D.EffectPassDescription desc = pass.Description;
            _name = desc.Name;
            _layoutMap = map;

            //Fetch annotations
            int annoCount = desc.AnnotationCount;
            _annotations = new D3D10EffectAnnotationCollection();
            for(int i = 0; i < annoCount; i++) {
                D3D.EffectVariable anno = _pass.GetAnnotationByIndex(i);
                D3D.EffectVariableDescription annoDesc = anno.Description;
                D3D.EffectVariable annoVar = anno.AsString();
                if(annoVar != null) {
                    _annotations.Add(new D3D10EffectAnnotation(annoVar, annoDesc.Name, annoDesc.Semantic, EffectParameterClass.Object, EffectParameterType.String));
                }
                annoVar = anno.AsScalar();
                if(annoVar != null) {
                    //Need to find a way to get the param type...
                    _annotations.Add(new D3D10EffectAnnotation(annoVar, annoDesc.Name, annoDesc.Semantic, EffectParameterClass.Scalar, EffectParameterType.Unknown));
                }
                annoVar = anno.AsMatrix();
                if(annoVar != null) {
                    _annotations.Add(new D3D10EffectAnnotation(annoVar, annoDesc.Name, annoDesc.Semantic, EffectParameterClass.Matrix, EffectParameterType.Single));
                }
            }
        }

        /// <summary>
        /// Creates a new instance of <see cref="D3D10EffectPass"/>.
        /// </summary>
        /// <param name="parent">The implementation parent.</param>
        /// <param name="effect">The D3D10 effect so we can query for variables.</param>
        /// <param name="pass">The D3D10 pass.</param>
        internal D3D10EffectPass(D3D10EffectImplementation parent, D3D.Effect effect, D3D.EffectPass pass) {
            _pass = pass;
            _renderer = parent.Renderer;

            D3D.EffectPassDescription desc = pass.Description;
            _name = desc.Name;
            _layoutMap = new InputLayoutMap(_renderer, desc.Signature);

            //Fetch annotations
            int annoCount = desc.AnnotationCount;
            _annotations = new D3D10EffectAnnotationCollection();
            for(int i = 0; i < annoCount; i++) {
                D3D.EffectVariable anno = _pass.GetAnnotationByIndex(i);
                D3D.EffectVariableDescription annoDesc = anno.Description;
                D3D.EffectVariable annoVar = anno.AsString();
                if(annoVar != null) {
                    _annotations.Add(new D3D10EffectAnnotation(annoVar, annoDesc.Name, annoDesc.Semantic, EffectParameterClass.Object, EffectParameterType.String));
                    continue;
                }
                annoVar = anno.AsScalar();
                if(annoVar != null) {
                    _annotations.Add(new D3D10EffectAnnotation(annoVar, annoDesc.Name, annoDesc.Semantic, EffectParameterClass.Scalar, EffectParameterType.Unknown));
                    continue;
                }
                annoVar = anno.AsMatrix();
                if(annoVar != null) {
                    _annotations.Add(new D3D10EffectAnnotation(annoVar, annoDesc.Name, annoDesc.Semantic, EffectParameterClass.Matrix, EffectParameterType.Unknown));
                    continue;
                }
                annoVar = anno.AsVector();
                if(annoVar != null) {
                    _annotations.Add(new D3D10EffectAnnotation(annoVar, annoDesc.Name, annoDesc.Semantic, EffectParameterClass.Vector, EffectParameterType.Unknown));
                }
            }

            //Now to reflect the parameters and start adding them to the shader
            D3D.EffectPassShaderDescription vertexShaderDesc = pass.VertexShaderDescription;
            D3D.EffectShaderVariable vertexShaderVar = vertexShaderDesc.Variable;
            ParseShader(parent, effect, vertexShaderVar);
            ParseShader(parent, effect, pass.PixelShaderDescription.Variable);

            //And the standard input layout
            CreateInputLayout(vertexShaderVar, vertexShaderDesc.Index);
        }

        private void CreateInputLayout(D3D.EffectShaderVariable shaderVar, int shaderIndex) {
            D3D.ShaderDescription shaderDesc;
            shaderVar.GetShaderDescription(shaderIndex, out shaderDesc);

            int paramCount = shaderDesc.InputParameterCount;

            //Dispose of the description's bytecode/sig...bytecode will be reported as a leak in ObjectTable
            shaderDesc.Bytecode.Dispose();
            shaderDesc.Signature.Dispose();

            VertexElement[] elements = new VertexElement[paramCount];
            D3D.ShaderParameterDescription desc;
            int offset = 0;
            for(int i = 0; i < paramCount; i++) {
                shaderVar.GetInputParameterDescription(shaderIndex, i, out desc);
                elements[i] = CreateVertexElement(desc, ref offset);
            }

            _layoutMap.CreateLayout(new VertexDeclaration(elements));
        }

        internal D3D10EffectPass Clone(D3D.EffectPass pass) {
            return new D3D10EffectPass(_renderer, pass, _layoutMap);
        }

        /// <summary>
        /// Applies the pass to the device. This sets up the rendering state before
        /// geometry rendering commences.
        /// </summary>
        public void Apply() {
            _pass.Apply();
            _renderer.OnApplyPass(this);
        }

        private static VertexElement CreateVertexElement(D3D.ShaderParameterDescription desc, ref int offset) {
            D3DC.RegisterComponentType compType = desc.ComponentType;
            D3DC.RegisterComponentMaskFlags components = desc.UsageMask;
            String semanticName = desc.SemanticName;
            uint semanticIndex = desc.SemanticIndex;

            if(D3D10Helper.IsFlagSet((int)components, (int) D3DC.RegisterComponentMaskFlags.All)) {
                if(semanticName.Equals("COLOR")) {
                    VertexElement element = new VertexElement(VertexSemantic.Color, (int) semanticIndex, VertexFormat.Color, offset);
                    offset += 4;
                    return element;
                } else {
                    VertexElement element = new VertexElement(D3D10Helper.FromD3DVertexSemantic(semanticName), (int) semanticIndex, VertexFormat.Vector4, offset);
                    offset += 16;
                    return element;
                }
            } else if(D3D10Helper.IsFlagSet((int) components, (int) (D3DC.RegisterComponentMaskFlags.ComponentX | D3DC.RegisterComponentMaskFlags.ComponentY | D3DC.RegisterComponentMaskFlags.ComponentZ))) {
                VertexElement element = new VertexElement(D3D10Helper.FromD3DVertexSemantic(semanticName), (int) semanticIndex, VertexFormat.Vector3, offset);
                offset += 12;
                return element;
            } else if(D3D10Helper.IsFlagSet((int) components, (int) (D3DC.RegisterComponentMaskFlags.ComponentX | D3DC.RegisterComponentMaskFlags.ComponentY))) {
                VertexElement element = new VertexElement(D3D10Helper.FromD3DVertexSemantic(semanticName), (int) semanticIndex, VertexFormat.Vector2, offset);
                offset += 8;
                return element;
            } else if(D3D10Helper.IsFlagSet((int) components, (int) (D3DC.RegisterComponentMaskFlags.ComponentX))) {
                VertexElement element = new VertexElement(D3D10Helper.FromD3DVertexSemantic(semanticName), (int) semanticIndex, VertexFormat.Single, offset);
                offset += 4;
                return element;
            }

            throw new ArgumentException("Invalid vertex element");
        }

        private static void ParseShader(D3D10EffectImplementation parent, D3D.Effect effect, D3D.EffectShaderVariable shaderVar) {
            D3D.EffectVariable testElem = shaderVar.GetElement(0);

            if(testElem == null) {
                if(shaderVar.IsValid) {
                    D3D.ShaderDescription shDesc;
                    shaderVar.GetShaderDescription(0, out shDesc);
                    D3DC.ShaderReflection reflect = new D3DC.ShaderReflection(shDesc.Bytecode);
                    ParseShaderParameters(parent, effect, reflect);

                    //Dispose of the reflection variable *and* the bytecode/signature of the shader desc...bytecode will reported as a leaked in ObjectTable
                    reflect.Dispose();
                    shDesc.Signature.Dispose();
                    shDesc.Bytecode.Dispose();
                }
            } else {
                int k = 1;
                while((testElem = shaderVar.GetElement(k)) != null) {
                    if(testElem.IsValid) {
                        D3D.ShaderDescription shDesc;
                        testElem.AsShader().GetShaderDescription(0, out shDesc);
                        D3DC.ShaderReflection reflect = new D3DC.ShaderReflection(shDesc.Bytecode);
                        ParseShaderParameters(parent, effect, reflect);

                        //Dispose of the reflection variable *and* the bytecode/signature of the shader desc...bytecode will reported as a leaked in ObjectTable
                        reflect.Dispose();
                        shDesc.Signature.Dispose();
                        shDesc.Bytecode.Dispose();
                    }
                    k++;
                }
            }
        }

        private static void ParseShaderParameters(D3D10EffectImplementation parent, D3D.Effect effect, D3DC.ShaderReflection reflect) {
            D3DC.ShaderDescription desc = reflect.Description;
            int constantBuffers = desc.ConstantBuffers;
            int boundResources = desc.BoundResources;

            //Grab parameters
            for(int i = 0; i < constantBuffers; i++) {
                D3DC.ConstantBuffer cb = reflect.GetConstantBuffer(i);
                for(int j = 0; j < cb.Description.Variables; j++) {
                    D3DC.ShaderReflectionVariable cbVar = cb.GetVariable(j);
                    String name = cbVar.Description.Name;
                    if(parent.Parameters[name] == null) {
                        D3D.EffectVariable param = effect.GetVariableByName(name);
                        D3DC.ShaderReflectionType paramType = cbVar.GetVariableType();
                        D3DC.ShaderTypeDescription typeDesc = paramType.Description;
                        if(Want(typeDesc.Class, typeDesc.Type)) {
                            ((D3D10EffectParameterCollection) parent.Parameters).Add(new D3D10EffectParameter(param, paramType));
                        }
                    } else {
                        break;
                    }
                }
            }
            //Grab resources
            for(int i = 0; i < boundResources; i++) {
                D3DC.InputBindingDescription inputBinding = reflect.GetResourceBindingDescription(i);
                String name = inputBinding.Name;
                if(parent.Parameters[name] == null) {
                    D3D.EffectVariable param = effect.GetVariableByName(name);
                    if(Want(inputBinding.Type, inputBinding.Dimension)) {
                        ((D3D10EffectParameterCollection) parent.Parameters).Add(new D3D10EffectParameter(param, inputBinding));
                    }
                }
            }
        }

        private static bool Want(D3DC.ShaderVariableClass paramClass, D3DC.ShaderVariableType paramType) {
            switch(paramClass) {
                case D3DC.ShaderVariableClass.Scalar:
                    switch(paramType) {
                        case D3DC.ShaderVariableType.Bool:
                        case D3DC.ShaderVariableType.Float:
                        case D3DC.ShaderVariableType.Int:
                            return true;
                        default:
                            return false;
                    }
                case D3DC.ShaderVariableClass.Vector:
                    return true;
                case D3DC.ShaderVariableClass.Struct:
                    switch(paramType) {
                        case D3DC.ShaderVariableType.Void:
                            return true;
                        default:
                            return false;
                    }
                case D3DC.ShaderVariableClass.MatrixColumns:
                case D3DC.ShaderVariableClass.MatrixRows:
                    return true;
                case D3DC.ShaderVariableClass.Object:
                    switch(paramType) {
                        case D3DC.ShaderVariableType.String:
                        case D3DC.ShaderVariableType.Texture:
                        case D3DC.ShaderVariableType.Texture1D:
                        case D3DC.ShaderVariableType.Texture2D:
                        case D3DC.ShaderVariableType.Texture3D:
                        case D3DC.ShaderVariableType.TextureCube:
                        case D3DC.ShaderVariableType.Texture1DArray:
                        case D3DC.ShaderVariableType.Texture2DArray:
                        case D3DC.ShaderVariableType.TextureCubeArray:
                            return true;
                        default:
                            return false;
                    }
                default:
                    return false;
            }
        }

        private static bool Want(D3DC.ShaderInputType inputType, SlimDX.Direct3D11.ShaderResourceViewDimension srv) {
            switch(inputType) {
                case D3DC.ShaderInputType.Texture:
                    switch(srv) {
                        case SlimDX.Direct3D11.ShaderResourceViewDimension.Texture1D:
                        case SlimDX.Direct3D11.ShaderResourceViewDimension.Texture2D:
                        case SlimDX.Direct3D11.ShaderResourceViewDimension.Texture3D:
                        case SlimDX.Direct3D11.ShaderResourceViewDimension.TextureCube:
                        case SlimDX.Direct3D11.ShaderResourceViewDimension.Texture1DArray:
                        case SlimDX.Direct3D11.ShaderResourceViewDimension.Texture2DArray:
                        case SlimDX.Direct3D11.ShaderResourceViewDimension.TextureCubeArray:
                            return true;
                        default:
                            return false;
                    }
                default:
                    return false;
            }
        }
    }
}
